import numpy as np
from sklearn.model_selection import train_test_split

def SSE(a,y):
    r = y-a
    return r.dot(r)

def SST(y):
    a = np.average(y)
    return SSE(a,y)

def attributeSelect(X,y):
        attribute = None
        threshold = None
        sst = None
        for a in range(X.shape[1]):
            x = np.array(X[:,a])
            x.sort()
            for i in range(1,X.shape[0]):
                if x[i-1] != x[i]:
                    t = (x[i-1]+x[i])/2
                    sst_ = SST(y[X[:,a]<=t]) + SST(y[X[:,a]>t])
                    if sst == None or sst_ < sst:
                        sst = sst_
                        attribute = a
                        threshold = t
        return attribute,threshold

class Node:
    def __init__(self):
        self.attribute = None
        self.threshold = None
        self.isLeaf = None
        self.samples_num = None
        self.predict_val = None
        self.left = None
        self.right = None
        return
    
    def decisionMaking(self,x):
        y_p = None
        if self.isLeaf:
            y_p = self.predict_val
        else:
            if x[self.attribute] <= self.threshold:
                y_p = self.left.decisionMaking(x)
            else:
                y_p = self.right.decisionMaking(x)
        return y_p

    def divide(self,X,y,maxdepth=None,depth=0):
        self.samples_num = y.size
        self.predict_val = np.average(y)
        if maxdepth == None or depth < maxdepth:
            a,t = attributeSelect(X,y)
            if a != None:
                self.isLeaf = False
                self.attribute = a
                self.threshold = t
                self.left = Node()
                self.left.divide(X[X[:,a]<=t,:],y[X[:,a]<=t],maxdepth,depth+1)
                self.right = Node()
                self.right.divide(X[X[:,a]>t,:],y[X[:,a]>t],maxdepth,depth+1)
            else:
                self.isLeaf = True
        else:
            self.isLeaf = True
        return

    def postPrun(self,X,y):
        if y.size == 0:
            sse = 0
        else:
            if self.isLeaf:
                sse = SSE(self.predict_val,y)
            else:
                sse1 = SSE(self.predict_val,y)
                left = X[:,self.attribute]<=self.threshold
                right = X[:,self.attribute]>self.threshold
                sse2 = self.left.postPrun(X[left,:],y[left]) + self.right.postPrun(X[right,:],y[right])
                if  sse1 <= sse2:
                    sse = sse1
                    self.isLeaf = True
                    self.left = None
                    self.right = None
                else:
                    sse = sse2
        return sse
    
    def prePrunDivide(self,X,y,X_,y_,maxdepth=None,depth=0):
        self.samples_num = y.size
        self.predict_val = np.average(y)
        if maxdepth == None or depth < maxdepth:
            a,t = attributeSelect(X,y)
            if a != None:
                left = X[:,a] <= t
                right = X[:,a] > t
                left_ = X_[:,a] <= t
                right_ = X_[:,a] > t
                sse1 = SSE(self.predict_val,y_)
                sse2 = SSE(np.average(y[left]),y_[left_]) + SSE(np.average(y[right]),y_[right_])
                if sse1 > sse2:
                    self.isLeaf = False
                    self.attribute = a
                    self.threshold = t
                    self.left = Node()
                    self.left.prePrunDivide(X[left,:],y[left],X_[left_],y_[left_],maxdepth,depth+1)
                    self.right = Node()
                    self.right.prePrunDivide(X[right,:],y[right],X_[right_],y_[right_],maxdepth,depth+1)
                else:
                    self.isLeaf = True
            else:
                self.isLeaf = True
        else:
            self.isLeaf = True
        return

    def show(self,fn,depth=0,trace=[]):
        for i in range(depth):
            if trace[i] == True:
                print("|   ",end='')
            else:
                print("    ",end='')
        print("|---",end='')
        print(fn[self.attribute]," <= %.2f"%self.threshold,", samples: ",self.left.samples_num,end='')
        trace.append(True)
        if(self.left.isLeaf):
            print(",","\033[1;31;36m","value: %.2f"%self.left.predict_val,"\033[0m")
        else:
            print('')
            self.left.show(fn,depth+1,trace)
        trace.pop()
        for i in range(depth):
            if trace[i] == True:
                print("|   ",end='')
            else:
                print("    ",end='')
        print("|---",end='')
        print(fn[self.attribute]," > %.2f"%self.threshold,", samples: ",self.right.samples_num,end='')
        trace.append(False)
        if(self.right.isLeaf):
            print(",","\033[1;31;36m","value: %.2f"%self.right.predict_val,"\033[0m")
        else:
            print('')
            self.right.show(fn,depth+1,trace)
        trace.pop()
        return
            

class RegressionTree:
    def __init__(self,prun=True,postprun=True,maxdepth=None):
        self.root = Node()
        self.prun = prun
        self.post = postprun
        self.maxdepth = maxdepth
        return
    
    def fit(self,X,y):
        X = np.array(X)
        y = np.array(y)
        if self.prun:
            X_train,X_test,y_train,y_test = train_test_split(X,y)
            if self.post:
                self.root.divide(X_train,y_train,self.maxdepth,0)
                self.root.postPrun(X_test,y_test)
            else:
                self.root.prePrunDivide(X_train,y_train,X_test,y_test,self.maxdepth,0)
        else:
            self.root.divide(X,y,self.maxdepth,0)
        return

    def predict(self,X):
        X = np.array(X)
        y_p = np.empty(X.shape[0])
        for i in range(X.shape[0]):
            y_p[i] = self.root.decisionMaking(X[i,:])
        return y_p
    
    def show(self,feature_names,target_names):
        print("samples: ",self.root.samples_num,", target: ",target_names[0])
        if(self.root.isLeaf == False):
            self.root.show(feature_names)
        print('')
        return
